from kuai.iotool import *
from kuai.strtool import *
from kuai.element import *
from kuai.mol import *
from kuai.strtool import extname
from cStringIO import StringIO

import os.path

class MoleculeReader(NamedReader):
    pass

class MolFileReader(MoleculeReader):
    BOND_ORDER = [SINGLE_BOND, DOUBLE_BOND, TRIPLE_BOND, PARTIAL_BOND]
    def read(self, file):
        file.next(); file.next(); file.next(); # Ignore 3 line
        line = file.next()
        natoms = int(line[0:3].strip())
        nbonds = int(line[3:6].strip())
        atoms = []
        for i in range(natoms):
            line = file.next()
            coords = XYZ(float(line[0:10].strip()), float(line[10:20].strip()), float(line[20:30].strip()))
            atoms.append(Atom(line[30,34].strip(), coords))
        
        bonds = []
        for i in range(nbonds):
            line = file.next()
            atom1 = line[0:3].strip()
            atom2 = line[3:6].strip()
            order = BOND_ORDER[int(line[6:9].strip())]
            bonds.append(Bond(atoms[atom1-1], atoms[atom2-1], order))
        return Molecule(atoms, bonds) 


def read_mol(file, type = None):
    if type is None:
        type = extname(file)
    from kuai.dff.dffmsd import reader
    file = open(file)
    result = reader.read(file)
    file.close()
    return result

def parse_mol(data, type):
    file = StringIO(data)
    from kuai.dff.dffmsd import reader
    return reader.read(file)

class MoleculeWriter(Writer):
    pass

def write_mol(file, mol, type = None):
    if type is None:
        type = extname(file)
    file = open(file, "w") 
    from kuai.dff.dffmsd import writer
    result = writer.write(file, mol)
    file.close()
    return result

def format_mol(mol, type = None):
    result = StringIO()
    from kuai.dff.dffmsd import writer
    writer.write(result, mol)
    return result.getvalue()